{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "  code by Tae Hwan Jung(Jeff Jung) @graykode, Derek Miller @dmmiller612, modify by wmathor\n",
    "  Reference : https://github.com/jadore801120/attention-is-all-you-need-pytorch\n",
    "              https://github.com/JayParks/transformer\n",
    "'''\n",
    "import math\n",
    "import torch\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as Data\n",
    "\n",
    "# S: Symbol that shows starting of decoding input\n",
    "# E: Symbol that shows starting of decoding output\n",
    "# P: Symbol that will fill in blank sequence if current batch data size is short than time steps\n",
    "sentences = [\n",
    "        # enc_input           dec_input         dec_output\n",
    "        ['ich mochte ein bier P', 'S i want a beer .', 'i want a beer . E'],\n",
    "        ['ich mochte ein cola P', 'S i want a coke .', 'i want a coke . E']\n",
    "]\n",
    "\n",
    "# Padding Should be Zero\n",
    "src_vocab = {'P' : 0, 'ich' : 1, 'mochte' : 2, 'ein' : 3, 'bier' : 4, 'cola' : 5}\n",
    "src_vocab_size = len(src_vocab)\n",
    "\n",
    "tgt_vocab = {'P' : 0, 'i' : 1, 'want' : 2, 'a' : 3, 'beer' : 4, 'coke' : 5, 'S' : 6, 'E' : 7, '.' : 8}\n",
    "idx2word = {i: w for i, w in enumerate(tgt_vocab)}\n",
    "tgt_vocab_size = len(tgt_vocab)\n",
    "\n",
    "src_len = 5 # enc_input max sequence length\n",
    "tgt_len = 6 # dec_input(=dec_output) max sequence length\n",
    "\n",
    "# Transformer Parameters\n",
    "d_model = 512  # Embedding Size\n",
    "d_ff = 2048 # FeedForward dimension\n",
    "d_k = d_v = 64  # dimension of K(=Q), V\n",
    "n_layers = 6  # number of Encoder of Decoder Layer\n",
    "n_heads = 8  # number of heads in Multi-Head Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "src_vocab, tgt_vocab, idx2word"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[1, 2, 3, 4, 0],\n",
       "         [1, 2, 3, 5, 0]]),\n",
       " tensor([[6, 1, 2, 3, 4, 8],\n",
       "         [6, 1, 2, 3, 5, 8]]),\n",
       " tensor([[1, 2, 3, 4, 8, 7],\n",
       "         [1, 2, 3, 5, 8, 7]]))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def make_data(sentences):\n",
    "    enc_inputs, dec_inputs, dec_outputs = [], [], []\n",
    "    for i in range(len(sentences)):\n",
    "      enc_input = [[src_vocab[n] for n in sentences[i][0].split()]] # [[1, 2, 3, 4, 0], [1, 2, 3, 5, 0]]\n",
    "      dec_input = [[tgt_vocab[n] for n in sentences[i][1].split()]] # [[6, 1, 2, 3, 4, 8], [6, 1, 2, 3, 5, 8]]\n",
    "      dec_output = [[tgt_vocab[n] for n in sentences[i][2].split()]] # [[1, 2, 3, 4, 8, 7], [1, 2, 3, 5, 8, 7]]\n",
    "\n",
    "      enc_inputs.extend(enc_input) # extend 方法用于将列表中的元素添加到另一个列表\n",
    "      dec_inputs.extend(dec_input)\n",
    "      dec_outputs.extend(dec_output)\n",
    "\n",
    "    return torch.LongTensor(enc_inputs), torch.LongTensor(dec_inputs), torch.LongTensor(dec_outputs)\n",
    "\n",
    "enc_inputs, dec_inputs, dec_outputs = make_data(sentences)\n",
    "enc_inputs, dec_inputs, dec_outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class MyDataSet(Data.Dataset):\n",
    "  def __init__(self, enc_inputs, dec_inputs, dec_outputs):\n",
    "    super(MyDataSet, self).__init__()\n",
    "    self.enc_inputs = enc_inputs\n",
    "    self.dec_inputs = dec_inputs\n",
    "    self.dec_outputs = dec_outputs\n",
    "  \n",
    "  def __len__(self):\n",
    "    return self.enc_inputs.shape[0]\n",
    "  \n",
    "  def __getitem__(self, idx):\n",
    "    return self.enc_inputs[idx], self.dec_inputs[idx], self.dec_outputs[idx]\n",
    "\n",
    "loader = Data.DataLoader(MyDataSet(enc_inputs, dec_inputs, dec_outputs), 2, True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 查看第一个样本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[1, 2, 3, 5, 0],\n",
       "         [1, 2, 3, 4, 0]]),\n",
       " tensor([[6, 1, 2, 3, 5, 8],\n",
       "         [6, 1, 2, 3, 4, 8]]),\n",
       " tensor([[1, 2, 3, 5, 8, 7],\n",
       "         [1, 2, 3, 4, 8, 7]])]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 获取第一个数据样本\n",
    "first_sample = next(iter(loader))\n",
    "first_sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 位置编码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, dropout=0.1, max_len=5000):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        # d_model 是词嵌入的维度\n",
    "        pe = torch.zeros(max_len, d_model) # pe.shape=(max_len, d_model)\n",
    "        # unsqueeze是用来升维的--position.shape=(max_len,1)\n",
    "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
    "        '''\n",
    "        position.shape=(max_len, 1)\n",
    "        dic_term.shape=(d_model/2)-->256\n",
    "        position * div_term --> (max_len, d_model/2)\n",
    "        '''\n",
    "        pe[:, 0::2] = torch.sin(position * div_term) # 这里用到了广播机制\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe = pe.unsqueeze(0).transpose(0, 1) # pe.shape=(max_len,1,512)\n",
    "        self.register_buffer('pe', pe)\n",
    "\n",
    "    def forward(self, x):\n",
    "        '''\n",
    "        x: [seq_len, batch_size, d_model]\n",
    "        '''\n",
    "        x = x + self.pe[:x.size(0), :]\n",
    "        return self.dropout(x)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pad Mask\n",
    "\n",
    "由于在 Encoder 和 Decoder 中都需要进行 mask 操作，因此就无法确定这个函数的参数中 `seq_len` 的值，如果是在 Encoder 中调用的，`seq_len` 就等于 `src_len`；如果是在 Decoder 中调用的，`seq_len` 就有可能等于 `src_len`，也有可能等于 `tgt_len`（因为 Decoder 有两次 mask）\n",
    "\n",
    "这个函数最核心的一句代码是 `seq_k.data.eq(0)`，这句的作用是返回一个大小和 `seq_k` 一样的 tensor，只不过里面的值只有 True 和 False。如果 `seq_k` 某个位置的值等于 0，那么对应位置就是 True，否则即为 False。举个例子，输入为 `seq_data = [1, 2, 3, 4, 0]`，`seq_data.data.eq(0)` 就会返回 `[False, False, False, False, True]`\n",
    "\n",
    "剩下的代码主要是扩展维度，强烈建议读者打印出来，看看最终返回的数据是什么样子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_attn_pad_mask(seq_q, seq_k):\n",
    "    '''\n",
    "    seq_q: [batch_size, seq_len]\n",
    "    seq_k: [batch_size, seq_len]\n",
    "    seq_len could be src_len or it could be tgt_len\n",
    "    seq_len in seq_q and seq_len in seq_k maybe not equal\n",
    "    '''\n",
    "    batch_size, len_q = seq_q.size()\n",
    "    batch_size, len_k = seq_k.size()\n",
    "    # eq(zero) is PAD token\n",
    "    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)  # [batch_size, 1, len_k], False is masked\n",
    "    return pad_attn_mask.expand(batch_size, len_q, len_k)  # [batch_size, len_q, len_k]\n",
    "\n",
    "# Subsequence Mask\n",
    "'''\n",
    "Subsequence Mask 只有 Decoder 会用到，主要作用是屏蔽未来时刻单词的信息。\n",
    "首先通过 np.ones() 生成一个全 1 的方阵，然后通过 np.triu() 生成一个上三角矩阵\n",
    "'''\n",
    "def get_attn_subsequence_mask(seq):\n",
    "    '''\n",
    "    seq: [batch_size, tgt_len]\n",
    "    '''\n",
    "    attn_shape = [seq.size(0), seq.size(1), seq.size(1)]\n",
    "    # k=1表明对角线上移一个单位，subsequence_mask 右上方都为1，左下方都为0，主对角线上是0\n",
    "    subsequence_mask = np.triu(np.ones(attn_shape), k=1) # Upper triangular matrix\n",
    "    # 将 subsequence_mask 的 NumPy 数组转换为 PyTorch 张量，并将其元素类型转换为 byte 类型\n",
    "    subsequence_mask = torch.from_numpy(subsequence_mask).byte()\n",
    "    return subsequence_mask # [batch_size, tgt_len, tgt_len]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 验证代码--草稿"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2, 6, 6]\n",
      "[[[0. 1. 1. 1. 1. 1.]\n",
      "  [0. 0. 1. 1. 1. 1.]\n",
      "  [0. 0. 0. 1. 1. 1.]\n",
      "  [0. 0. 0. 0. 1. 1.]\n",
      "  [0. 0. 0. 0. 0. 1.]\n",
      "  [0. 0. 0. 0. 0. 0.]]\n",
      "\n",
      " [[0. 1. 1. 1. 1. 1.]\n",
      "  [0. 0. 1. 1. 1. 1.]\n",
      "  [0. 0. 0. 1. 1. 1.]\n",
      "  [0. 0. 0. 0. 1. 1.]\n",
      "  [0. 0. 0. 0. 0. 1.]\n",
      "  [0. 0. 0. 0. 0. 0.]]]\n",
      "torch.Size([2, 6, 6])\n",
      "tensor([[[0, 1, 1, 1, 1, 1],\n",
      "         [0, 0, 1, 1, 1, 1],\n",
      "         [0, 0, 0, 1, 1, 1],\n",
      "         [0, 0, 0, 0, 1, 1],\n",
      "         [0, 0, 0, 0, 0, 1],\n",
      "         [0, 0, 0, 0, 0, 0]],\n",
      "\n",
      "        [[0, 1, 1, 1, 1, 1],\n",
      "         [0, 0, 1, 1, 1, 1],\n",
      "         [0, 0, 0, 1, 1, 1],\n",
      "         [0, 0, 0, 0, 1, 1],\n",
      "         [0, 0, 0, 0, 0, 1],\n",
      "         [0, 0, 0, 0, 0, 0]]], dtype=torch.uint8)\n"
     ]
    }
   ],
   "source": [
    "q = torch.arange(0,12).reshape(2, 6)\n",
    "k = torch.arange(0,10).reshape(2, 5)\n",
    "mask =  get_attn_pad_mask(seq_q = q, seq_k=k)\n",
    "# print(q, k, mask.shape, mask)\n",
    "\n",
    "attn_shape = [q.size(0), q.size(1), q.size(1)]\n",
    "print(attn_shape)\n",
    "subsequence_mask = np.triu(np.ones(attn_shape), k=1) # Upper triangular matrix\n",
    "print(subsequence_mask)\n",
    "subsequence_mask = torch.from_numpy(subsequence_mask).byte()\n",
    "print(subsequence_mask.shape)\n",
    "print(subsequence_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 验证广播机制"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "position = torch.tensor([[1],\n",
    "                         [2],\n",
    "                         [3],\n",
    "                         [4],\n",
    "                         [5]])\n",
    "div_term = torch.tensor([1,2,3,4,5,6,7,8])\n",
    "position * div_term"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 对transpose函数的理解"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pe = torch.tensor([[1, 2, 3, 4],\n",
    "              [1, 1, 1, 1],\n",
    "              [2, 2, 2, 2],\n",
    "              [3, 4, 5, 6],\n",
    "              [2, 1, 5, 9],\n",
    "              [9, 2, 10, 3]])\n",
    "print(pe.shape) #(6,4)\n",
    "print(pe.transpose(0, 1)) #shape=(4,6)\n",
    "pe = pe.unsqueeze(0).transpose(0, 1) #shape=(6,1,4)\n",
    "print(pe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ScaledDotProductAttention\n",
    "\n",
    "这里要做的是，通过 Q 和 K 计算出 scores，然后将 scores 和 V 相乘，得到每个单词的 context vector\n",
    "\n",
    "第一步是将 Q 和 K 的转置相乘没什么好说的，相乘之后得到的 scores 还不能立刻进行 softmax，需要和 attn_mask 相加，把一些需要屏蔽的信息屏蔽掉，attn_mask 是一个仅由 True 和 False 组成的 tensor，并且一定会保证 attn_mask 和 scores 的维度四个值相同（不然无法做对应位置相加）\n",
    "\n",
    "mask 完了之后，就可以对 scores 进行 softmax 了。然后再与 V 相乘，得到 context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ScaledDotProductAttention(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ScaledDotProductAttention, self).__init__()\n",
    "\n",
    "    def forward(self, Q, K, V, attn_mask):\n",
    "        '''\n",
    "        Q: [batch_size, n_heads, len_q, d_k]\n",
    "        K: [batch_size, n_heads, len_k, d_k]\n",
    "        V: [batch_size, n_heads, len_v(=len_k), d_v]\n",
    "        attn_mask: [batch_size, n_heads, seq_len, seq_len]\n",
    "        '''\n",
    "        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k) # scores : [batch_size, n_heads, len_q, len_k]\n",
    "        # 上三角的部分被置为-inf\n",
    "        scores.masked_fill_(attn_mask, -1e9) # Fills elements of self tensor with value where mask is True.\n",
    "        \n",
    "        attn = nn.Softmax(dim=-1)(scores)\n",
    "        context = torch.matmul(attn, V) # [batch_size, n_heads, len_q, d_v]\n",
    "        return context, attn\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False)\n",
    "        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=False)\n",
    "        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=False)\n",
    "        self.fc = nn.Linear(n_heads * d_v, d_model, bias=False)\n",
    "    def forward(self, input_Q, input_K, input_V, attn_mask):\n",
    "        '''\n",
    "        input_Q: [batch_size, len_q, d_model]\n",
    "        input_K: [batch_size, len_k, d_model]\n",
    "        input_V: [batch_size, len_v(=len_k), d_model]\n",
    "        attn_mask: [batch_size, seq_len, seq_len]\n",
    "        '''\n",
    "        residual, batch_size = input_Q, input_Q.size(0)\n",
    "        # (B, S, D) -proj-> (B, S, D_new) -split-> (B, S, H, W) -trans-> (B, H, S, W)\n",
    "        Q = self.W_Q(input_Q).view(batch_size, -1, n_heads, d_k).transpose(1,2)  # Q: [batch_size, n_heads, len_q, d_k]\n",
    "        K = self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1,2)  # K: [batch_size, n_heads, len_k, d_k]\n",
    "        V = self.W_V(input_V).view(batch_size, -1, n_heads, d_v).transpose(1,2)  # V: [batch_size, n_heads, len_v(=len_k), d_v]\n",
    "\n",
    "        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) # attn_mask : [batch_size, n_heads, seq_len, seq_len]\n",
    "\n",
    "        # context: [batch_size, n_heads, len_q, d_v], attn: [batch_size, n_heads, len_q, len_k]\n",
    "        context, attn = ScaledDotProductAttention()(Q, K, V, attn_mask)\n",
    "        context = context.transpose(1, 2).reshape(batch_size, -1, n_heads * d_v) # context: [batch_size, len_q, n_heads * d_v]\n",
    "        output = self.fc(context) # [batch_size, len_q, d_model]\n",
    "        return nn.LayerNorm(d_model).cuda()(output + residual), attn\n",
    "\n",
    "class PoswiseFeedForwardNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(PoswiseFeedForwardNet, self).__init__()\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(d_model, d_ff, bias=False),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(d_ff, d_model, bias=False)\n",
    "        )\n",
    "    def forward(self, inputs):\n",
    "        '''\n",
    "        inputs: [batch_size, seq_len, d_model]\n",
    "        '''\n",
    "        residual = inputs\n",
    "        output = self.fc(inputs)\n",
    "        return nn.LayerNorm(d_model).cuda()(output + residual) # [batch_size, seq_len, d_model]\n",
    "\n",
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.enc_self_attn = MultiHeadAttention()\n",
    "        self.pos_ffn = PoswiseFeedForwardNet()\n",
    "\n",
    "    def forward(self, enc_inputs, enc_self_attn_mask):\n",
    "        '''\n",
    "        enc_inputs: [batch_size, src_len, d_model]\n",
    "        enc_self_attn_mask: [batch_size, src_len, src_len]\n",
    "        '''\n",
    "        # enc_outputs: [batch_size, src_len, d_model], attn: [batch_size, n_heads, src_len, src_len]\n",
    "        enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask) # enc_inputs to same Q,K,V\n",
    "        enc_outputs = self.pos_ffn(enc_outputs) # enc_outputs: [batch_size, src_len, d_model]\n",
    "        return enc_outputs, attn\n",
    "\n",
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.dec_self_attn = MultiHeadAttention()\n",
    "        self.dec_enc_attn = MultiHeadAttention()\n",
    "        self.pos_ffn = PoswiseFeedForwardNet()\n",
    "\n",
    "    def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask):\n",
    "        '''\n",
    "        dec_inputs: [batch_size, tgt_len, d_model]\n",
    "        enc_outputs: [batch_size, src_len, d_model]\n",
    "        dec_self_attn_mask: [batch_size, tgt_len, tgt_len]\n",
    "        dec_enc_attn_mask: [batch_size, tgt_len, src_len]\n",
    "        '''\n",
    "        # dec_outputs: [batch_size, tgt_len, d_model], dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len]\n",
    "        dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs, dec_self_attn_mask)\n",
    "        # dec_outputs: [batch_size, tgt_len, d_model], dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]\n",
    "        dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs, enc_outputs, dec_enc_attn_mask)\n",
    "        dec_outputs = self.pos_ffn(dec_outputs) # [batch_size, tgt_len, d_model]\n",
    "        return dec_outputs, dec_self_attn, dec_enc_attn\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.src_emb = nn.Embedding(src_vocab_size, d_model)\n",
    "        self.pos_emb = PositionalEncoding(d_model)\n",
    "        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])\n",
    "\n",
    "    def forward(self, enc_inputs):\n",
    "        '''\n",
    "        enc_inputs: [batch_size, src_len]\n",
    "        '''\n",
    "        enc_outputs = self.src_emb(enc_inputs) # [batch_size, src_len, d_model]\n",
    "        enc_outputs = self.pos_emb(enc_outputs.transpose(0, 1)).transpose(0, 1) # [batch_size, src_len, d_model]\n",
    "        enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs) # [batch_size, src_len, src_len]\n",
    "        enc_self_attns = []\n",
    "        for layer in self.layers:\n",
    "            # enc_outputs: [batch_size, src_len, d_model], enc_self_attn: [batch_size, n_heads, src_len, src_len]\n",
    "            enc_outputs, enc_self_attn = layer(enc_outputs, enc_self_attn_mask)\n",
    "            enc_self_attns.append(enc_self_attn)\n",
    "        return enc_outputs, enc_self_attns\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model)\n",
    "        self.pos_emb = PositionalEncoding(d_model)\n",
    "        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])\n",
    "\n",
    "    def forward(self, dec_inputs, enc_inputs, enc_outputs):\n",
    "        '''\n",
    "        dec_inputs: [batch_size, tgt_len]\n",
    "        enc_intpus: [batch_size, src_len]\n",
    "        enc_outputs: [batsh_size, src_len, d_model]\n",
    "        '''\n",
    "        dec_outputs = self.tgt_emb(dec_inputs) # [batch_size, tgt_len, d_model]\n",
    "        dec_outputs = self.pos_emb(dec_outputs.transpose(0, 1)).transpose(0, 1).cuda() # [batch_size, tgt_len, d_model]\n",
    "        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs).cuda() # [batch_size, tgt_len, tgt_len]\n",
    "        dec_self_attn_subsequence_mask = get_attn_subsequence_mask(dec_inputs).cuda() # [batch_size, tgt_len, tgt_len]\n",
    "        dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + dec_self_attn_subsequence_mask), 0).cuda() # [batch_size, tgt_len, tgt_len]\n",
    "\n",
    "        dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs) # [batc_size, tgt_len, src_len]\n",
    "\n",
    "        dec_self_attns, dec_enc_attns = [], []\n",
    "        for layer in self.layers:\n",
    "            # dec_outputs: [batch_size, tgt_len, d_model], dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len], dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]\n",
    "            dec_outputs, dec_self_attn, dec_enc_attn = layer(dec_outputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask)\n",
    "            dec_self_attns.append(dec_self_attn)\n",
    "            dec_enc_attns.append(dec_enc_attn)\n",
    "        return dec_outputs, dec_self_attns, dec_enc_attns\n",
    "\n",
    "class Transformer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Transformer, self).__init__()\n",
    "        self.encoder = Encoder().cuda()\n",
    "        self.decoder = Decoder().cuda()\n",
    "        self.projection = nn.Linear(d_model, tgt_vocab_size, bias=False).cuda()\n",
    "    def forward(self, enc_inputs, dec_inputs):\n",
    "        '''\n",
    "        enc_inputs: [batch_size, src_len]\n",
    "        dec_inputs: [batch_size, tgt_len]\n",
    "        '''\n",
    "        # tensor to store decoder outputs\n",
    "        # outputs = torch.zeros(batch_size, tgt_len, tgt_vocab_size).to(self.device)\n",
    "        \n",
    "        # enc_outputs: [batch_size, src_len, d_model], enc_self_attns: [n_layers, batch_size, n_heads, src_len, src_len]\n",
    "        enc_outputs, enc_self_attns = self.encoder(enc_inputs)\n",
    "        # dec_outpus: [batch_size, tgt_len, d_model], dec_self_attns: [n_layers, batch_size, n_heads, tgt_len, tgt_len], dec_enc_attn: [n_layers, batch_size, tgt_len, src_len]\n",
    "        dec_outputs, dec_self_attns, dec_enc_attns = self.decoder(dec_inputs, enc_inputs, enc_outputs)\n",
    "        dec_logits = self.projection(dec_outputs) # dec_logits: [batch_size, tgt_len, tgt_vocab_size]\n",
    "        return dec_logits.view(-1, dec_logits.size(-1)), enc_self_attns, dec_self_attns, dec_enc_attns\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "model = Transformer().cuda()\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=0)\n",
    "optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.99)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0001 loss = 0.064609\n",
      "Epoch: 0002 loss = 0.029503\n",
      "Epoch: 0003 loss = 0.021982\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(3):\n",
    "    for enc_inputs, dec_inputs, dec_outputs in loader:\n",
    "      '''\n",
    "      enc_inputs: [batch_size, src_len]\n",
    "      dec_inputs: [batch_size, tgt_len]\n",
    "      dec_outputs: [batch_size, tgt_len]\n",
    "      '''\n",
    "      enc_inputs, dec_inputs, dec_outputs = enc_inputs.cuda(), dec_inputs.cuda(), dec_outputs.cuda()\n",
    "      # outputs: [batch_size * tgt_len, tgt_vocab_size]\n",
    "      outputs, enc_self_attns, dec_self_attns, dec_enc_attns = model(enc_inputs, dec_inputs)\n",
    "      loss = criterion(outputs, dec_outputs.view(-1))\n",
    "      print('Epoch:', '%04d' % (epoch + 1), 'loss =', '{:.6f}'.format(loss))\n",
    "\n",
    "      optimizer.zero_grad()\n",
    "      loss.backward()\n",
    "      optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def greedy_decoder(model, enc_input, start_symbol):\n",
    "    \"\"\"\n",
    "    For simplicity, a Greedy Decoder is Beam search when K=1. This is necessary for inference as we don't know the\n",
    "    target sequence input. Therefore we try to generate the target input word by word, then feed it into the transformer.\n",
    "    Starting Reference: http://nlp.seas.harvard.edu/2018/04/03/attention.html#greedy-decoding\n",
    "    :param model: Transformer Model\n",
    "    :param enc_input: The encoder input\n",
    "    :param start_symbol: The start symbol. In this example it is 'S' which corresponds to index 4\n",
    "    :return: The target input\n",
    "    \"\"\"\n",
    "    enc_outputs, enc_self_attns = model.encoder(enc_input)\n",
    "    dec_input = torch.zeros(1, 0).type_as(enc_input.data)\n",
    "    terminal = False\n",
    "    next_symbol = start_symbol\n",
    "    while not terminal:         \n",
    "        dec_input = torch.cat([dec_input.detach(),torch.tensor([[next_symbol]],dtype=enc_input.dtype).cuda()],-1)\n",
    "        dec_outputs, _, _ = model.decoder(dec_input, enc_input, enc_outputs)\n",
    "        projected = model.projection(dec_outputs)\n",
    "        prob = projected.squeeze(0).max(dim=-1, keepdim=False)[1]\n",
    "        next_word = prob.data[-1]\n",
    "        next_symbol = next_word\n",
    "        if next_symbol == tgt_vocab[\".\"]:\n",
    "            terminal = True\n",
    "        print(next_word)            \n",
    "    return dec_input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1, device='cuda:0')\n",
      "tensor(2, device='cuda:0')\n",
      "tensor(3, device='cuda:0')\n",
      "tensor(4, device='cuda:0')\n",
      "tensor(8, device='cuda:0')\n",
      "tensor([1, 2, 3, 4, 0], device='cuda:0') -> ['i', 'want', 'a', 'beer', '.']\n",
      "tensor(1, device='cuda:0')\n",
      "tensor(2, device='cuda:0')\n",
      "tensor(3, device='cuda:0')\n",
      "tensor(5, device='cuda:0')\n",
      "tensor(8, device='cuda:0')\n",
      "tensor([1, 2, 3, 5, 0], device='cuda:0') -> ['i', 'want', 'a', 'coke', '.']\n"
     ]
    }
   ],
   "source": [
    "# Test\n",
    "enc_inputs, _, _ = next(iter(loader))\n",
    "enc_inputs = enc_inputs.cuda()\n",
    "for i in range(len(enc_inputs)):\n",
    "    greedy_dec_input = greedy_decoder(model, enc_inputs[i].view(1, -1), start_symbol=tgt_vocab[\"S\"])\n",
    "    predict, _, _, _ = model(enc_inputs[i].view(1, -1), greedy_dec_input)\n",
    "    predict = predict.data.max(1, keepdim=True)[1]\n",
    "    print(enc_inputs[i], '->', [idx2word[n.item()] for n in predict.squeeze()])\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "newd2l",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
